from functools import wraps
import inspect
import time
import os

import numpy as np


class QuantileKeeper:
    def __init__(self, quantile=0.95, T=8000) -> None:
        self.quantile = quantile
        self.revert_quantile = 1 - quantile
        self.T = T
        self.t = 0
        self.buffer_size = int(T * self.revert_quantile) + 100
        self.buffer = [None] * (self.buffer_size + 1)
        self.buffer_i = 0

    def put(self, v):
        for i in range(self.buffer_i, -1, -1):
            if (i > 0) and (self.buffer[i - 1] < v):
                self.buffer[i] = self.buffer[i - 1]
            else:
                self.buffer[i] = v
                break
        self.buffer_i = min(self.buffer_i + 1, self.buffer_size)
        self.t += 1

    def get(self):
        p = int(self.t * self.revert_quantile)
        return self.buffer[p]

    def get_neighbor(self, distance):
        p = int(self.t * self.revert_quantile) - distance
        p = np.clip(p, a_min=0, a_max=self.buffer_i-1)
        return self.buffer[p]


class QuantileUpdater:
    """
    实时获取、更新0.95分位点所在时刻，并且可以获取0.95分位点时刻的相邻时刻
    """
    def __init__(self, value_list, quantile_index) -> None:
        self.value_list = value_list
        self.sorted_index = np.argsort(value_list)
        self.rank_index = np.zeros_like(self.sorted_index)
        self.rank_index[self.sorted_index] = np.arange(len(value_list))

        self.T = len(value_list)
        self.quantile_index = quantile_index

    def get(self, offset=0):
        t = self.sorted_index[self.quantile_index + offset]
        v = self.value_list[t]
        return t, v

    def get_t(self, offset=0):
        return self.sorted_index[self.quantile_index + offset]

    def get_v(self, offset=0):
        return self.value_list[self.get_t(offset)]

    def get_all(self):
        t_arr = self.sorted_index[self.quantile_index:]
        v_arr = self.value_list[t_arr]
        return t_arr, v_arr

    def update(self, t, new_v):
        old_v = self.value_list[t]
        if old_v == new_v:
            return
        old_rank = self.rank_index[t]

        if new_v > old_v:
            _si = old_rank + 1
            while (_si < self.T) and (new_v > self.value_list[self.sorted_index[_si]]):
                self.rank_index[self.sorted_index[_si]] -= 1
                self.rank_index[t] += 1
                self.sorted_index[_si - 1] = self.sorted_index[_si]
                _si += 1
            self.sorted_index[_si - 1] = t
        else:
            _si = old_rank - 1
            while (_si >= 0) and (new_v < self.value_list[self.sorted_index[_si]]):
                self.rank_index[self.sorted_index[_si]] += 1
                self.rank_index[t] -= 1
                self.sorted_index[_si + 1] = self.sorted_index[_si]
                _si -= 1
            self.sorted_index[_si + 1] = t
        self.value_list[t] = new_v


class LargestKeeper:
    """
    维护最大流量排序，可以实时获取最大流量的值及其所在时刻，并且如果流量没有被分配完，可以重新放入该排序对象中
    使用分捅的方式实现， 维护桶内有序
    """
    def __init__(self, value_list, bucket_num=100) -> None:
        self.T = len(value_list)
        self.bucket_num = min(self.T, bucket_num)

        self.value_t_list = list(zip(value_list, range(self.T)))
        self.range_split_list = None
        self.bucket_list = None
        self.bucket_inner_i = 0  # 最大bucket数组内，遍历到了第几个元素
        self._init_bucket()

    def _init_bucket(self):
        self.value_t_list.sort()
        self.range_split_list = [0] * self.bucket_num
        self.bucket_list = [None] * self.bucket_num
        split_points = np.linspace(0, self.T, self.bucket_num + 1, dtype=np.int64)
        for k in range(self.bucket_num):
            i = split_points[k]
            j = split_points[k + 1]
            bucket = self.value_t_list[i:j]
            self.bucket_list[k] = bucket
            self.range_split_list[k] = bucket[-1][0]
        self.range_split_list.pop()
        self.bucket_inner_i = len(self.bucket_list[-1]) - 1

    def get(self):
        if not self.bucket_list:
            return 0, None

        return self.bucket_list[-1][self.bucket_inner_i]

    def next(self):
        if not self.bucket_list:
            return False

        self.bucket_inner_i -= 1
        while self.bucket_inner_i < 0:
            self.bucket_list.pop()
            if not self.bucket_list:
                return False
            else:
                self.range_split_list.pop()
                bucket = self.bucket_list[-1]
                bucket.sort()
                self.bucket_inner_i = len(bucket) - 1
        return True

    def add(self, new_v, new_t):
        if not self.bucket_list:
            bucket = [(new_v, new_t)]
            self.bucket_list.append(bucket)
            self.bucket_inner_i = 0
            return

        new_bucket_i = np.searchsorted(self.range_split_list, new_v)
        new_tuple = (new_v, new_t)
        if new_bucket_i == (len(self.bucket_list) - 1):  # 正好等于当前正在遍历的桶
            bucket = self.bucket_list[-1]
            bucket = bucket[:self.bucket_inner_i + 1]
            bucket.append(new_tuple)
            bucket.sort()
            self.bucket_list[-1] = bucket
            self.bucket_inner_i += 1
        else:
            self.bucket_list[new_bucket_i].append(new_tuple)

    def __repr__(self) -> str:
        return "(%s,%s)" % self.get()


def sort_plot(X, jpg_name=None):
    import matplotlib.pyplot as plt
    from constants import OUTPUT_DIR
    import time
    jpg_file = f"{OUTPUT_DIR}/{time.strftime('%Y%m%d_%H%M%S', time.localtime(time.time()))}_{jpg_name}.jpg"

    X = np.array(X)
    p = int(np.ceil(0.95 * X.shape[0]) - 1)
    if len(X.shape) == 1:
        sorted_x = sorted(X)

        plt.figure()
        plt.plot(sorted_x)
        plt.scatter([p], [sorted_x[p]], c='r', label='*')
        if jpg_name:
            print(f"save plot to: {jpg_file}")
            plt.savefig(jpg_file, dpi=300)
        else:
            plt.show()
    else:
        N = X.shape[1]
        X = np.sort(X, axis=0)
        cols = np.argsort(X[p, :])

        fig, axes_list = plt.subplots(N, 1, figsize=(10, 10 * N))
        for i, c in enumerate(cols):
            axes = axes_list[i]
            axes.set_title(f"Site:{c}")
            axes.plot(X[:, c])
            axes.scatter([p], [X[p, c]], c='r', label='*')
        if jpg_name:
            print(f"save plot to: {jpg_file}")
            fig.savefig(jpg_file)
        else:
            plt.show()


def assign_plot(t, all_assign_matrix, jpg_name=None):
    import matplotlib.pyplot as plt
    import time

    from constants import OUTPUT_DIR
    from data_io import M, N, site_names, client_names, demand_matrix, site_bandwidth_arr, reachable_matrix, quantile_index

    assign_matrix = all_assign_matrix[t]
    demand_arr = demand_matrix[t]
    site_wb_devoted = assign_matrix.sum(axis=0)
    site_wb_remained = site_bandwidth_arr - site_wb_devoted
    site_quantile_arr = np.sort(all_assign_matrix.sum(axis=1), axis=0)[quantile_index]

    plt.figure(figsize=(25, 5))
    plt.axis('off')
    max_size = max(M, N)
    x_scale = 12
    plt.axis([-5, max_size * x_scale + 5, -0.2, 1.2])

    client_x_arr = np.linspace(0, max_size, num=M) * x_scale
    client_y_arr = np.zeros_like(client_x_arr)
    plt.scatter(client_x_arr, client_y_arr, s=90, c='b')
    for client_id, client_name in enumerate(client_names):
        plt.text(client_x_arr[client_id] - 8,
                 client_y_arr[client_id] - 0.11,
                 f"{client_name}\nD:{demand_arr[client_id]}",
                 fontsize=8,
                 rotation=45)

    site_x_arr = np.linspace(0, max_size, num=N) * x_scale
    site_y_arr = np.ones_like(site_x_arr)

    for site_id, site_name in enumerate(site_names):
        _c = 'k'
        if not site_wb_remained[site_id]:
            _c = 'r'
        else:
            if site_wb_devoted[site_id] >= site_quantile_arr[site_id]:
                _c = 'y'
            else:
                _c = 'navy'

        plt.scatter([site_x_arr[site_id]],
                    [site_y_arr[site_id]],
                    s=90,
                    c=_c)
        plt.text(site_x_arr[site_id] - 8,
                 site_y_arr[site_id],
                 f"{site_name}\n"
                 f"Up:{site_bandwidth_arr[site_id]}\n"
                 f"Re:{site_wb_devoted[site_id]}\n"
                 f"Rm:{site_wb_remained[site_id]}\n"
                 f"Q:{site_quantile_arr[site_id]}",
                 fontsize=8,
                 rotation=-45)

    for client_id, site_id in zip(*np.where(reachable_matrix)):
        plt.plot([client_x_arr[client_id], site_x_arr[site_id]], [client_y_arr[client_id], site_y_arr[site_id]], color='lightgrey')

    for client_id, site_id in zip(*np.where(assign_matrix)):
        x1 = client_x_arr[client_id]
        x2 = site_x_arr[site_id]
        y1 = client_y_arr[client_id]
        y2 = site_y_arr[site_id]
        plt.plot([x1, x2], [y1, y2], color='g')
        _lambda = np.random.random()
        plt.text(x1 + (x2 - x1) * _lambda,
                 y1 + (y2 - y1) * _lambda,
                 f"{assign_matrix[client_id, site_id]}",
                 fontsize=8)

    if jpg_name:
        jpg_file = f"{OUTPUT_DIR}/{time.strftime('%Y%m%d_%H%M%S', time.localtime(time.time()))}_{jpg_name}.jpg"
        print(f"save plot to: {jpg_file}")
        plt.savefig(jpg_file)
    else:
        plt.show()


def mf(file_path):  # mkdir for file_path
    _dir = os.path.dirname(file_path)
    _dir and os.makedirs(_dir, exist_ok=True)
    return file_path


def md(dirname):  # mkdir at dirname
    dirname and os.makedirs(dirname, exist_ok=True)
    return dirname


def save_pkl(obj, file_path):
    import pickle
    with open(mf(file_path), "wb") as _f:
        pickle.dump(obj, _f)


def load_pkl(file_path, default=None):
    import pickle
    if os.path.exists(file_path):
        with open(file_path, "rb") as _f:
            return pickle.load(_f)
    elif default is not None:
        return default
    else:
        raise FileNotFoundError(f"not such file : {file_path}")


def eval(assigner):
    import time
    t1 = time.time()
    for t in range(assigner.T):
        if not (t % 1000):
            print(f"t: {t}, current_score: {assigner.current_score()}")
        assigner.next_assign()
    t2 = time.time()
    print(f"time cost: {t2 - t1} seconds")
    print(f"score: {assigner.current_score()}")


def logging_time_cost(log_input=True, log_input_keys=None, log_output=False):
    """
    :param log_input:是否打印input
    :param log_input_keys:字符串数组，指明打印哪些kwargs中的参数；默认None表示打印全部input
    :param log_output:是否打印output
    :param log_level:日志级别，默认DEBUG
    :return:
    """

    def _wraper(func):
        _func = func
        while hasattr(_func, '__wrapped__'):
            _func = _func.__wrapped__
        _varnames_ = _func.__code__.co_varnames[:_func.__code__.co_argcount]
        _log_input_keys = log_input_keys or _varnames_

        @wraps(func)
        def decorated(*args, **kwargs):
            func_params = inspect.getcallargs(_func, *args, **kwargs)
            if log_input:
                _input_log_str = ', '.join(map(lambda k: f"{k}={func_params[k]}", _log_input_keys))
            else:
                _input_log_str = ''
            print(f"start func: {func.__name__}({_input_log_str})")
            t1 = time.time()
            ans = func(*args, **kwargs)
            t2 = time.time()
            print(f"finish func: {func.__name__}, time cost: {round(t2 - t1, 3)}s")
            if log_output:
                print(f"output is : {ans}")
            return ans

        return decorated

    return _wraper
